package com.will.nio.selector;

import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.util.Iterator;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicInteger;

import static com.will.nio.buffer.ByteBufferUtil.debugAll;

/**
 * @author Will
 * @description:
 *              利用多线程优化事件处理
 * @date 2023-04-27 22:20
 * @version: 1.0
 */
@Slf4j
public class MultiThreadServer {
    public static void main(String[] args) throws IOException {
        new Boss().bossInit();
    }

    /**
     * 单线程（主线程）专门处理 accept 事件
     */
    static class Boss {
        public void bossInit() throws IOException {
            ServerSocketChannel ssc = ServerSocketChannel.open().bind(new InetSocketAddress(8080));
            ssc.configureBlocking(false);
            Selector selector = Selector.open();
            ssc.register(selector, SelectionKey.OP_ACCEPT);
            // 创建固定数量的 worker 并初始化
            Worker[] workers = new Worker[Runtime.getRuntime().availableProcessors()];
            for (int i = 0; i < workers.length; i++) {
                workers[i] = new Worker("worker-" + i);
            }
            AtomicInteger index = new AtomicInteger();
            while (true) {
                selector.select();
                Iterator<SelectionKey> iterator = selector.selectedKeys().iterator();
                while (iterator.hasNext()) {
                    SelectionKey key = iterator.next();
                    if (key.isAcceptable()) {
                        SocketChannel sc = ssc.accept();
                        sc.configureBlocking(false);
                        log.debug("read before...");
                        // round robin 轮询
                        workers[index.getAndIncrement() % workers.length].initWorker(sc);
                        log.debug("read after...");
                    }
                    iterator.remove();
                }
            }
        }
    }

    /**
     * 专门处理 read 事件
     */
    static class Worker implements Runnable {
        private Thread thread;
        private Selector selector;
        private String name;
        private volatile boolean isInit = false;

        private ConcurrentLinkedQueue<Runnable> tasks = new ConcurrentLinkedQueue<>();

        public Worker(String name) {
            this.name = name;
        }

        /**
         * 多个 Worker，但每个 Worker 只需要初始化一次
         * @param sc
         * @throws IOException
         */
        public void initWorker(SocketChannel sc) throws IOException {
            if (!isInit) {
                this.selector = Selector.open();
                this.thread = new Thread(this, name);
                this.thread.start();
                this.isInit = true;
            }
            // Worker 初始化后，selector 调用了 select 方法，处于阻塞模式，
            // 此时 channel 注册到 selector 上会失败
            // 因此需要借助阻塞队列，让 channel 成功注册到 selector 上
            tasks.add(() -> {
                try {
                    sc.register(this.selector, SelectionKey.OP_READ);
                } catch (ClosedChannelException e) {
                    e.printStackTrace();
                }
            });
            // 让 selector 恢复运行
            this.selector.wakeup();
        }

        @Override
        public void run() {
            while (true) {
                try {
                    selector.select();
                    // 将 channel 注册到 selector 上，并关注 read 事件
                    Runnable task = tasks.poll();
                    if (task != null) {
                        task.run();
                    }
                    Iterator<SelectionKey> iterator = selector.selectedKeys().iterator();
                    while (iterator.hasNext()) {
                        SelectionKey key = iterator.next();
                        if (key.isReadable()) {
                            try {
                                SocketChannel sc = (SocketChannel) key.channel();
                                ByteBuffer buffer = ByteBuffer.allocate(16);
                                int read = sc.read(buffer);
                                if (read == -1) {
                                    key.cancel();
                                } else {
                                    log.debug("当前线程：{}", Thread.currentThread().getName());
                                    buffer.flip();
                                    debugAll(buffer);
                                }
                            } catch (IOException e) {
                                key.cancel();
                                e.printStackTrace();
                            }
                        }
                        iterator.remove();
                    }
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }
}
